import StringIO
from django.http import HttpResponse
from django.conf import settings
# Note that for performance -- and hopefully bug-elimination --
# reasons, all matplotlib-related imports have been pushed down into
# the methods that use them.


def plot_results(results, run, plottype='final', cmap_idx=0, max_range=100, standalone=False,
                 lonmin=None, latmin=None, lonmax=None, latmax=None):
    """If optional arg `standalone' is True, will draw in graticules,
    colourbar, coastline, etc.  max_range as a percentage can be used
    to provide more resolution on the final-positions plots where the
    values are all small and hard to discern."""
    try:
        stats = results.finalStats if plottype == 'final' else results.stats

        # check if a restricted area is requested:
        try:
            def _doround(val, inc): return round(val/inc)*inc
            lonmin = _doround(max(float(lonmin), results.lon_min) if lonmin else results.lon_min,
                              results.lon_inc)
            latmin = _doround(max(float(latmin), results.lat_min) if latmin else results.lat_min,
                              results.lat_inc)
            lonmax = _doround(min(float(lonmax), results.lon_max) if lonmax else results.lon_max,
                              results.lon_inc)
            latmax = _doround(min(float(latmax), results.lat_max) if latmax else results.lat_max,
                              results.lat_inc)
            assert(lonmin < lonmax)
            assert(latmin < latmax)

            # all ok, so see if we need to trim any (note we need to
            # work in array coords here, not geo):
            lonmindelta = (lonmin - results.lon_min) / results.lon_inc
            latmindelta = (latmin - results.lat_min) / results.lat_inc
            lonmaxdelta = (results.lon_max - lonmax) / results.lon_inc
            latmaxdelta = (results.lat_max - latmax) / results.lat_inc
            if any((lonmindelta, latmindelta, lonmaxdelta, latmaxdelta)):
                rows,cols = stats.shape
                stats = stats[latmindelta:(rows-latmaxdelta),
                              lonmindelta:(cols-lonmaxdelta)]
        except:
            # reset to normal values if there's anything funny
            # (including, eg, float parse errors):
            lonmin = results.lon_min
            latmin = results.lat_min
            lonmax = results.lon_max
            latmax = results.lat_max

        import numpy as np
        # using the difference between max and min seems to lead to occasional off-by-one errors:
        rows,cols = stats.shape
        xs = np.arange(lonmin, lonmin+cols*results.lon_inc, results.lon_inc)
        ys = np.arange(latmin, latmin+rows*results.lat_inc, results.lat_inc)

        from matplotlib.figure import Figure, figaspect
        fig = Figure()
        if standalone:
            graph = fig.add_axes([0.1,0.1, 0.9, 0.9])
            resolution = 'i'
            fontsize='x-small'
        else:
            graph = fig.add_axes([0,0,1,1])
            graph.axis('off')
            resolution = None

        from mpl_toolkits.basemap import Basemap
        from matplotlib.cm import get_cmap
        from matplotlib.colors import Normalize
        from matplotlib.pyplot import close
        from matplotlib.patches import Circle, Rectangle
        m = Basemap(ax=graph, llcrnrlon=lonmin, llcrnrlat=latmin, urcrnrlon=lonmax, urcrnrlat=latmax, resolution=resolution)
        if standalone:
            from math import floor, ceil
            m.drawcoastlines()
            m.drawmeridians(np.arange(ceil(lonmin), floor(lonmax+1)), labels=[0,0,0,1], fontsize=fontsize)
            m.drawparallels(np.arange(ceil(latmin), floor(latmax+1)), labels=[1,0,0,0], fontsize=fontsize)
        width, height = figaspect((latmax-latmin)/(lonmax-lonmin))
        fig.set_size_inches((4*width, 4*height))

        xs, ys = m(xs, ys)
        vmax = float(max_range)/100;
        plot = m.pcolor(xs, ys, stats, norm=Normalize(vmin=0, vmax=vmax, clip=True), edgecolors='None', cmap=get_cmap(settings.COLOUR_MAPS[cmap_idx]))

        if standalone:
            fig.suptitle(str(run))
            cbax = fig.colorbar(plot, shrink=0.8).ax
            labels = cbax.get_yticklabels()
            for label in labels:
                label.set_fontsize(fontsize)

        # now draw the initial positions in:
        for poly in run.regions:
            xmin,ymin,xmax,ymax = poly.extent
            if run.isSource:
                patch = Rectangle(m(xmin,ymin), xmax-xmin, ymax-ymin, edgecolor='black', facecolor='None', aa=False)
            else:
                xradius, yradius = (xmax - xmin)/2, (ymax - ymin)/2
                xcentre,ycentre = xmin + xradius, ymin + yradius
                radius = min(xradius, yradius)
                patch = Circle(m(xcentre, ycentre), radius=radius, edgecolor='black', facecolor='None')
            graph.add_patch(patch)

        from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
        canvas = FigureCanvas(fig)
        fig.set_canvas(canvas)
        buf = StringIO.StringIO()
        fig.savefig(buf, format='png', transparent=True)
        close(fig)
        img = buf.getvalue()
        buf.close()
        return img
    except:
        import traceback, sys
        traceback.print_exc(); sys.stderr.flush()

def create_vert(fig, cmap_name):
    from matplotlib import colorbar
    from matplotlib.cm import get_cmap
    graph = fig.add_axes([0,0.01,0.9,0.98])
    fig.set_size_inches(0.25, 5)
    colorbar.ColorbarBase(graph, cmap=get_cmap(cmap_name), orientation='vertical')
def create_horizontal(fig, cmap_name):
    from matplotlib import colorbar
    from matplotlib.cm import get_cmap
    from matplotlib.figure import Figure, figaspect
    graph = fig.add_axes([0.01,0.1,0.98,0.9])
    fig.set_size_inches(figaspect(1.0/20))
    colorbar.ColorbarBase(graph, cmap=get_cmap(cmap_name), orientation='horizontal')
def colourbar(request):
    cmap_idx = request.GET.get('cmap', 0)
    try:
        cmap_idx = int(cmap_idx)
    except:
        cmap_idx = 0
    if cmap_idx < 0 or cmap_idx >= len(settings.COLOUR_MAPS):
        cmap_idx = 0
    cmap = settings.COLOUR_MAPS[cmap_idx]

    response = HttpResponse(content_type='image/png')
    # note: not using per-view cache here because that (afaik) ignores skips requests with get parameters:
    from django.core.cache import cache
    img = cache.get('cmap%d' % cmap_idx)
    if not img:
        from matplotlib.pyplot import close
        from matplotlib.figure import Figure
        from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
        buf = StringIO.StringIO()
        fig = Figure()
        create_vert(fig, cmap)
        canvas = FigureCanvas(fig)
        fig.set_canvas(canvas)
        fig.savefig(buf, format='png', transparent=True)
        close(fig)
        img = buf.getvalue()
        buf.close()
        cache.set('cmap%d' % cmap_idx, img, 24*60*60) # cache for a day

    response.write(img)
    return response

